package com.xiaojie.ratelimit;

import jdk.jfr.events.ExceptionThrownEvent;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.serializer.Jackson2JsonRedisSerializer;
import org.springframework.data.redis.serializer.RedisSerializer;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;

import javax.annotation.Resource;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * 使用计数器方式进行限流
 *
 * @author wangye
 * @classname RedisCountRateLimit
 * @date 2021/3/31 13:40
 **/
@Component
public class RedisRateLimit {

    DefaultRedisScript<Long> rs = new DefaultRedisScript<Long>();
    @Resource
    private RedisTemplate<String, Long> redisTemplate;
    private Jackson2JsonRedisSerializer argsSerializer = new Jackson2JsonRedisSerializer(Integer.class);
    private Jackson2JsonRedisSerializer argsSerializer2 = new Jackson2JsonRedisSerializer(Double.class);
    private Jackson2JsonRedisSerializer resultSerializer = new Jackson2JsonRedisSerializer(Long.class);

    /**
     * 计数限流
     *
     * @param maxlimit
     * @return
     */
    public Long call(int maxlimit) {
        rs.setScriptSource(new ResourceScriptSource(new ClassPathResource("ratelimit/countratelimit.lua")));
        rs.setResultType(Long.class);
        Long count = redisTemplate.execute(rs, argsSerializer, resultSerializer, Stream.of("countratelimit").collect(Collectors.toList()), maxlimit);
        return count;
    }

    /**
     * 令牌桶限流
     * 通过令牌产生的速率避免集中请求的临界问题
     * @param maxlimit 最大令牌数量
     * @param permits  获取令牌数
     * @param rate 每毫秒产生令牌数量
     * @param curTime 当前时间
     * @param bucketExpire 令牌过期时间
     * @param tokenInitNum 令牌初始化数量
     * @return
     */
    public Long bucketlimit(Long maxlimit, Long permits, Double rate, Long curTime,Long bucketExpire,Long tokenInitNum) {
        rs.setScriptSource(new ResourceScriptSource(new ClassPathResource("ratelimit/bucketratelimit.lua")));
        rs.setResultType(Long.class);
        Long count = redisTemplate.execute(rs,argsSerializer2,resultSerializer,
                Stream.of("bucket","last_mill_request_key ").collect(Collectors.toList()), maxlimit,permits,rate,curTime,bucketExpire,tokenInitNum);
        return count;
    }
}
